#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : species_count_accumulation.py
@Author  : Bing Liang
@Date    : 2025/10/15
@Description : 根据 RPM 矩阵计算累积物种数表
"""

import pandas as pd
from collections import defaultdict
from argparse import ArgumentParser, Namespace
from pathlib import Path


def main(args: Namespace):
    """
    主函数：计算物种累积数并输出结果
    """
    input_path = Path(args.input).resolve()
    output_path = Path(args.output).resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # 读取 RPM 文件
    df = pd.read_csv(input_path, sep="\t")

    # 计算每个物种在多少个样本中出现（RPM > 0）
    df["count"] = (df.iloc[:, 1:] > 0).sum(axis=1)

    # 按出现次数排序
    df = df.sort_values(by="count", ascending=False)

    # 输出累积物种数统计表
    with open(output_path, "w", encoding="utf-8") as fw:
        fw.write("SampleN\tSpeciesN\n")
        for i in range(1, df.shape[1]):  # 从1到样本总数
            species_count = df.query("count >= @i").shape[0]
            fw.write(f"{i}\t{species_count}\n")

    print(f"✅ 输出完成: {output_path}")


if __name__ == "__main__":
    parser = ArgumentParser(description="根据 RPM 矩阵计算累积物种数表")
    parser.add_argument("-i", "--input", required=True, help="输入 RPM 矩阵文件 (all.RPM.tsv)")
    parser.add_argument("-o", "--output", required=True, help="输出累积物种数文件 (species_counts.xls)")
    main(parser.parse_args())
